#!/usr/env/bin python
# -*- coding: utf-8 -*-

# @Time    : 2020-07-15 2:24 下午
# @Author  : yangdy
# @File    : 对角线遍历
# @Software: PyCharmCE
from typing import List


class Solution:
    def findDiagonalOrder(self, matrix: List[List[int]]) -> List[int]:
        ans = []
        if not matrix:
            return ans
        n = len(matrix)
        m = len(matrix[0])
        ss = [
            (-1, 1), (0, 1), (1, 0),
            (1, -1), (1, 0), (0, 1),
        ]
        x, y, i = 1, -1, 0
        lans = 0
        while lans < n * m:
            xx, yy = x + ss[i][0], y + ss[i][1]
            while 0 <= xx < n and 0 <= yy < m:
                x, y = xx, yy
                ans.append(matrix[x][y])
                lans += 1
                xx, yy = x + ss[i][0], y + ss[i][1]
            if lans >= n * m:
                break
            while not(0 <= xx < n and 0 <= yy < m):
                i = (i + 1) % 6
                xx, yy = x + ss[i][0], y + ss[i][1]
            x, y = xx, yy
            ans.append(matrix[x][y])
            lans += 1
            if i < 3:
                i = 3
            else:
                i = 0
        return ans


solve = Solution()

questions = [
    ([[1, 2, 3], [4, 5, 6], [7, 8, 9]],),
    ([[1]],),
]

for qi, q in enumerate(questions):
    print(f'#{qi}:', q)
    print('ans:', solve.findDiagonalOrder(q[0]))
